"""腾讯混元大模型独立实现类"""
import os
import json
import time
import logging
import requests
from typing import Any, Dict, List, Optional

class TencentHunyuanModel:
    """腾讯混元大模型独立实现类（支持重试机制）"""
    
    API_URL = "https://api.hunyuan.cloud.tencent.com/v1/chat/completions"
    
    def __init__(self, api_key: str, model: str = "hunyuan-turbos-latest", 
                 max_retries: int = 3, retry_delay: float = 2.0):
        """
        初始化腾讯混元模型
        
        :param api_key: API密钥
        :param model: 模型ID（默认hunyuan-turbos-latest）
        :param max_retries: 最大重试次数（默认3次）
        :param retry_delay: 重试延迟时间（秒，默认2秒）
        """
        self.api_key = api_key
        self.model = model
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.logger = logging.getLogger("TencentHunyuanModel")
        self._setup_logger()
    
    def _setup_logger(self):
        """配置日志记录器"""
        self.logger.setLevel(logging.INFO)
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.info("腾讯混元模型初始化完成 | 模型: %s | 最大重试次数: %d", 
                         self.model, self.max_retries)
    
    def chat_completions(self, messages: List[Dict], **kwargs) -> Dict[str, Any]:
        """
        发送聊天补全请求（带重试机制）
        
        :param messages: 消息列表
        :param kwargs: 可选参数
        :return: API响应字典
        """
        # 获取模型参数
        model = kwargs.get("model", self.model)
        max_tokens = kwargs.get("max_tokens", 1024)
        temperature = kwargs.get("temperature", 0.7)
        enable_enhancement = kwargs.get("enable_enhancement", True)
        timeout = kwargs.get("timeout", 30.0)
        
        # 构建请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        # 构建请求体
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "enable_enhancement": enable_enhancement
        }
        
        # 添加可选参数
        if "top_p" in kwargs:
            payload["top_p"] = kwargs["top_p"]
        if "stop_sequences" in kwargs:
            payload["stop_sequences"] = kwargs["stop_sequences"]
        
        # 重试机制
        for attempt in range(self.max_retries + 1):
            try:
                # 发送请求
                response = requests.post(
                    self.API_URL,
                    headers=headers,
                    json=payload,
                    timeout=timeout
                )
                
                # 处理HTTP错误
                if response.status_code != 200:
                    error_msg = f"API错误: {response.status_code} - {response.text}"
                    self.logger.error(error_msg)
                    return {"error": error_msg}
                
                return response.json()
            
            except requests.exceptions.Timeout:
                if attempt < self.max_retries:
                    # 指数退避策略
                    wait_time = self.retry_delay * (2 ** attempt)
                    self.logger.warning("请求超时，将在 %.1f 秒后重试（尝试 %d/%d）", 
                                        wait_time, attempt + 1, self.max_retries)
                    time.sleep(wait_time)
                else:
                    error_msg = "请求超时：在多次重试后仍失败"
                    self.logger.error(error_msg)
                    return {"error": error_msg}
                    
            except requests.exceptions.RequestException as e:
                error_msg = f"网络请求异常: {str(e)}"
                self.logger.error(error_msg)
                return {"error": error_msg}
        
        # 正常情况下不会执行到这里
        return {"error": "请求失败，重试次数用尽"}
    
    def text_chat(self, prompt: str, **kwargs) -> str:
        """
        执行纯文本对话
        
        :param prompt: 文本提示
        :param kwargs: 可选参数
        :return: 模型生成的文本响应
        """
        messages = [
            {
                "role": "user",
                "content": prompt
            }
        ]
        
        # 设置默认超时
        timeout = kwargs.pop('timeout', 30.0)
        response = self.chat_completions(messages, timeout=timeout, **kwargs)
        return self._parse_response(response)
    
    def _parse_response(self, response_data: Dict[str, Any]) -> str:
        """解析腾讯混元API响应"""
        if "error" in response_data:
            return response_data["error"]
        
        try:
            if "choices" in response_data and response_data["choices"]:
                return response_data["choices"][0]["message"]["content"].strip()
            
            # 记录异常响应
            self.logger.warning("响应格式异常: %s", json.dumps(response_data, indent=2))
            return "腾讯混元返回了无法解析的响应格式"
        
        except Exception as e:
            self.logger.exception("响应解析异常: %s", str(e))
            return f"响应解析失败: {str(e)}"
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取腾讯混元模型信息"""
        return {
            "model_name": "tencent_hunyuan",
            "api_endpoint": self.API_URL,
            "api_key": self.api_key,  # 注意：实际应用中可能需要隐藏敏感信息
            "model_id": self.model
        }

# 使用示例
if __name__ == "__main__":
    # 初始化模型（使用环境变量中的API密钥）
    api_key = os.getenv("HUNYUAN_API_KEY", "your_api_key_here")
    model = TencentHunyuanModel(
        api_key=api_key,
        model="hunyuan-turbos-latest",
        max_retries=5,
        retry_delay=1.5
    )
    
    # 测试文本对话
    response = model.text_chat(
        "你好，请介绍一下你自己",
        timeout=45.0,
        max_tokens=300,
        temperature=0.7
    )
    
    print("腾讯混元响应:")
    print(response)
    
    # 获取模型信息
    print("\n模型信息:")
    print(json.dumps(model.get_model_info(), indent=2))